package tests;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Random;

import stdp_rules.LogSTDP;
import utils.StateCamera;
import utils.StateCamera.CameraType;
import utils.math_utils.Coordinate;
import utils.math_utils.MathUtils;
import utils.math_utils.RandStream;
import Jama.Matrix;

import components.AbstractNeuron;
import components.MLNeuron;
import components.Network;
import components.Synapse;

import connectors.GridConnect;
import connectors.PreferrentialAttachment;

public class AdderTest {
	
	public enum Gate {
		OR, XOR, FA_SUM, FA_CARRY;
		
		public String getName(){
			String name;
			switch(this) {
				case OR : name = "Or"; break;
				case XOR : name = "XOR"; break;
				case FA_SUM : name = "Full Adder: Sum"; break;
				case FA_CARRY : name = "Full Adder: Carry"; break;
				default : name = null; break;
			}
			return name;
		}
		
		
		public boolean function (boolean [] ins) {
			switch(this) {
				case OR : return ins[0] || ins[1];
				case XOR : return ins[0] != ins[1];
				case FA_SUM : return fullAdderSum(ins[0], ins[1], ins[2]);
				case FA_CARRY : return fullAdderCarry(ins[0], ins[1], ins[2]);
				default : throw new IllegalArgumentException("No such gate.");
			}
		}
		
		private boolean fullAdderSum (boolean a, boolean b, boolean c) {
			return (a != b) != c;	
		}
		
		private boolean fullAdderCarry(boolean a, boolean b, boolean c) {
			return ((a != b) && c) || (a && b);
		}
		
	}

	public static String testFileName;
	public static final int firstPass = 1999;
	public static final int trainEpochs = 2000;
	public static final int testEpochs = 1000;
	public static final double timeStep = 0.1; //ms
	public static double time = 0;
	public static final double pulseDuration = 0.3;//ms
	public static final int epoch = 30; //ms	
	public static final int numRes = 1024;
	public static final int memTc = 30;
	private static final double [] StimCurrent = {5000, 5000, 5000}; //nA
	
	private StateCamera camera;
	private static int [] resDims = {16, 16, 4};
	private final Network net = new Network();
	private static boolean [][] trainingIn;
	private static boolean [][] testingIn;
	private static boolean [][] trainingTargs = new boolean [trainEpochs][1];
	private static boolean [][] testingTargs = new boolean [testEpochs][1];
	private MLNeuron [] res = new MLNeuron[numRes];
	private static int stimGrpSize = resDims[2] * 25;
	private MLNeuron [] stimGrp1 = new MLNeuron[stimGrpSize];
	private MLNeuron [] stimGrp2 = new MLNeuron[stimGrpSize];
	private MLNeuron [] stimGrp3 = new MLNeuron[stimGrpSize];
	private static byte [][] stateMatrix;
	private static Matrix wOut;
//	private static Matrix [] filters = new Matrix[epoch];

	
	public AdderTest(){
		initialize();
	}
	
	public void generateInputSets(Gate gate){
		if(gate == Gate.OR || gate == Gate.XOR) {
			trainingIn = new boolean[trainEpochs][2];
			testingIn = new boolean[testEpochs][2];
		} else {
			trainingIn = new boolean[trainEpochs][3];
			testingIn = new boolean[testEpochs][3];
		}
		Random rand = new Random();
		for(int i = 0; i < trainEpochs; i++) {
			for(int j = 0; j < trainingIn[i].length; j++) {
				trainingIn[i][j] = rand.nextBoolean();
			}
			trainingTargs[i][0] = gate.function(trainingIn[i]);
			printToFile("TrTargets", MathUtils.toDouble(trainingTargs));
			printToFile("TrIns", MathUtils.toDouble(trainingIn));
		}
		
		for(int i = 0; i < testEpochs; i++) {
			for(int j = 0; j < testingIn[i].length; j++) {
				testingIn[i][j] = rand.nextBoolean();
			}
			testingTargs[i][0] = gate.function(testingIn[i]);
		}
		
	}
	
	/**
	 * 
	 */
	private void initialize() {
		
		camera = new StateCamera(CameraType.Bin,
				timeStep, 30);			
		
//		res =(Network.initializeNeurons(numRes,
//				memTc, 15, 2.5, 0.2, camera));
		
		res = new MLNeuron[1024];
		
		boolean ex;
		for(int i = 0; i < 1024; i++) {
			ex = Math.random() > 0.2;
			res[i] = new MLNeuron(ex, i);
			
//			res[i].setMemPotential(RandStream.nextUniform(-80, 40));
//			res[i].setC(RandStream.nextGaussian(20, 5, 5, 35, true));
			
//			if(Math.random() > 0.5){
//				res[i].setV3(2);
//				res[i].setV4(30);
//				res[i].setG_Ca(4.4);
//				res[i].setPhi(0.04);
////				if(Math.random() > 0.5) {
////					res[i].setEpsilon(0.002);
////					
////				}
//			}
			
		}
	
		net.cubicLayout(new Coordinate(0, 0, 0), resDims,
				res);
		
		int x1 = 4;
		int x2 = 8;
		int x3 = 12;
		int ys = 6;	
		
		HashMap<Coordinate, AbstractNeuron> map = net.getNeuronMap();
		int ind1 = 0;
		int ind2 = 0;
		int ind3 = 0;
		
		for(int z = 0; z < resDims[2]; z++) {
			
			for(int x = x1 - 2; x <= x1 + 2; x++){
				for(int y = ys - 2; y <= ys + 2; y++){
					stimGrp1[ind1++] = (MLNeuron) map.get(new Coordinate(x, y, z));
//					 System.out.println(map.get(new Coordinate(x, y, z)).getIndex());
				}
			}
			
			for(int x = x2 - 2; x <= x2 + 2; x++){
				for(int y = ys - 2; y <= ys + 2; y++){
					stimGrp2[ind2++] = (MLNeuron) map.get(new Coordinate(x, y, z));
//					 System.out.println(map.get(new Coordinate(x, y, z)).getIndex());
				}
			}

			for(int x = x3 - 2; x <= x3 + 2; x++){
				for(int y = ys - 2; y <= ys + 2; y++){
					stimGrp3[ind3++] = (MLNeuron) map.get(new Coordinate(x, y, z));
//					System.out.println(map.get(new Coordinate(x, y, z)).getIndex());
				}
			}
			
		}	
		
		
		PreferrentialAttachment pCon = new PreferrentialAttachment();
		GridConnect gCon = new GridConnect();
		Network.initializeSynRandomizers(gCon);
		Network.initializeSynRandomizers(pCon);
		gCon.stochasticConnect(res, res, 3.0);
		pCon.addPreferentialConnections(res);
		pCon.addPreferentialConnections(res);
		pCon.addPreferentialConnections(res);
		pCon.addPreferentialConnections(res);
		double c = 0.0;
		int indx = 0;
		double [][] syns = new double[1024][1024];
		for(MLNeuron n : res){	
			c += n.getFanIn().size();
			for(Synapse s: n.getFanIn()) {
				syns[indx][s.getSrc().getIndex()] = s.getStrength();
			}
			indx++;
		}
		System.out.println("New Sparsity: " + c/(res.length * res.length));
		
        	                   // tcP,  tcD,  lRate,    nv,   alp,   bet,    cp,     wc 
		LogSTDP	stdp = new LogSTDP(15,   30,    0.1,  0.6f,  10.0f,  100f,  1.0f,  100f);
		for(MLNeuron n : res) {
			n.setStdpOn(false);
			n.setStdp(stdp);
			n.setReporter(true);
			n.setRecorder(camera);
			//Homeostatic plasticity ON
			//n.setHmP(false);
		}
		printToFile("SynapsesBefore.dat", syns);
	}
	
	
	/**
	 * Runs the adder test for cases where states are to be linearly combined.
	 * @param test
	 * @param inputs
	 * @param recording
	 */
	public static void runSeparate(AdderTest test, boolean [][] inputs,
			boolean recording, boolean testing){
		int num;
		String process;
		if(!recording && !testing){
			num = firstPass;
			process = "LC First Pass";
		} else if (recording && !testing) {
			num = trainEpochs;
			process = "LC Training";
		} else {
			num = testEpochs;
			process = "LC Testing";
		}
		System.out.println(process);
		StateCamera camera = test.getCamera();
		LinkedList<MLNeuron> spikingNeurons;
		boolean current;
		double per = 0.1;
		double runningPer = 0;
		double counter = 0;
		stateMatrix = new byte [num][numRes];
		
		for(int i = 0; i < inputs[0].length; i++) {
			if(recording)
				camera.readyFilm(inputs.length, numRes);
			for(int j = 0; j < num; j++) {
				current = inputs[j][i];
				MLNeuron [] stimGrp = null;
				if(current) {
					stimGrp = test.getStimGrp(i + 1);
//					for(MLNeuron n : stimGrp){
//						n.setI_bg(n.getI_bg() + StimCurrent[0]);
//					}
				}				
				int c = 0;
				double t = 0;
				while(t < epoch) {
					
					if(c < 3 && current){
						for(MLNeuron n : stimGrp){
							n.setI_bg(n.getI_bg() + StimCurrent[c]);
						}
					}
					
					
					if(c == 3 && stimGrp != null){
						for(MLNeuron n : stimGrp) {
							n.setI_bg(45);
						}
					}
					
					spikingNeurons = new LinkedList<MLNeuron>();
					for(MLNeuron n : test.getRes()) {
						if(n.partialUpdate(timeStep)){
							spikingNeurons.add(n);
						}						
						
					}
					
					while(!spikingNeurons.isEmpty()){
						spikingNeurons.poll().spike(timeStep);
					}
					
					if(recording)
						camera.collectFrame();
					
					t += timeStep;
					time += timeStep;
					c++;
				}
				
				runningPer = counter / (inputs.length * inputs[0].length);
				if(runningPer >= per) {
					per += 0.1;
					System.out.print(Math.floor(runningPer * 100) + "% ");
				}
				counter++;
			}
			if(recording) {
				MathUtils.arrayAddInPlace2D(stateMatrix,
						MathUtils.toByte(camera.getCapturedStates()));
			}
		}
		
		if(recording){	
			if(testing) {
				test.test(stateMatrix);
			} else {
				test.train(stateMatrix);
			}
		}		
	}
	
	
	/**
	 * 
	 * @param test
	 * @param inputs
	 * @param recording
	 */
	public static void run (AdderTest test, boolean [][] inputs,
			boolean recording, boolean testing) {
		
		int num;
		String process;
		if(!recording && !testing){
			num = firstPass;
			process = "First Pass";
		} else if (recording && !testing) {
			num = trainEpochs;
			process = "Training";
		} else {
			num = testEpochs;
			process = "Testing";
		}
		System.out.println(process);
		
		stateMatrix  = new byte[num][numRes];		
		StateCamera camera = test.getCamera();
		LinkedList<MLNeuron> spikingNeurons;
		boolean [] current;
		double per = 0.01;
		double runningPer = 0;
		try {
			PrintWriter printer = null;
			if(num == firstPass){
				FileWriter fw = new FileWriter("MemRes.dat");
				printer = new PrintWriter(fw);
			}
//			if(recording)
				camera.readyFilm(inputs.length, numRes);
			ArrayList<MLNeuron> stims = new ArrayList<MLNeuron>();
			for(int i = 0, m = inputs.length; i < m; i++) {
				current = inputs[i];
				for(int j = 0; j < inputs[i].length; j++) {
					if(current[j]){
						for(MLNeuron n : test.getStimGrp(j+1)){
							stims.add(n);
						}
					}
				}
				
				int c = 0;
				double t = 0;
				while(t < epoch) {
					
					if(c < 3 && !stims.isEmpty()) {
						for(MLNeuron n : stims) {
							n.setI_bg(StimCurrent[c]);
						}
					}
					
					if(c == 3 && !stims.isEmpty()){
						for(MLNeuron n : stims) {
							n.setI_bg(35);
						}
					}
					
					spikingNeurons = new LinkedList<MLNeuron>();
					for(MLNeuron n : test.getRes()) {
						if(n.partialUpdate(timeStep)){
							spikingNeurons.add(n);
						}
//						if(num == firstPass && (t == Math.floor(t))) {
//							printer.print((n.getIndex() + 1) + " ");
//							printer.print(n.getMemResistance());	
//							printer.println();					
//						}
					}
					
					while(!spikingNeurons.isEmpty()){
						spikingNeurons.poll().spike(timeStep);
					}
					
//					if(recording)
						camera.collectFrame();
					
					if(i == 200 && !recording && t >= 25 && t <= 25.5) {
						boolean [][] g = camera.getFrameStack();
						printToFile("EpSnapshot" + t +".dat", MathUtils.toDouble(g));
					}
					
					t += timeStep;
					time += timeStep;
					c++;
				}
				
				runningPer = (double)i/m;
				if(runningPer >= per) {
					per += 0.01;
					System.out.print(Math.floor(runningPer * 100) + "% ");
				}
				

				
			}
			
			if(recording){
				stateMatrix = MathUtils.toByte(camera.getCapturedStates());
				
			}
				if(!recording){
					printToFile("FPSM.dat", camera.getCapturedStates());
				}
			
			if(num == firstPass)
				printer.close();
		} catch(IOException ie) {
			ie.printStackTrace();
			System.exit(1);
		}
		if(recording){	
			if(testing) {
				test.test(stateMatrix);
			} else {
				test.train(stateMatrix);
			}
		}	
		System.out.println();
	}
	
	/**
	 * TRAINNG
	 * @param states
	 * @param inputs
	 */
	public void train(byte [][] states) {
		System.out.println();
		Matrix S;
		Matrix dOuts = new Matrix(MathUtils.toDouble(trainingTargs));
		double [][] sm;
			
		sm = MathUtils.toDouble(states);
			
		printToFile("TrainStates.dat", sm);
		MathUtils.addNoise(new RandStream(-0.0001, 0.0001), sm);
		S = new Matrix(sm);
		try {
			wOut = S.solve(dOuts);
			System.out.print("SUCCESS");
		} catch (RuntimeException re) {
			System.out.print("FAILURE");
			System.exit(1);
		}
		System.out.println();			
	}
	
	/**
	 * TEST
	 * @param states
	 * @param inputs
	 */
	public void test(byte [][] states) {
		Matrix S;
		double [][] dOuts = MathUtils.toDouble(testingTargs);
		double [][] sm;
		try {
			FileWriter fw = new FileWriter(testFileName);
			PrintWriter pw = new PrintWriter(fw);
			System.out.println();
			System.out.print("[");
			sm = MathUtils.toDouble(states);
			S = new Matrix(sm);
			double [][] outs = (S.times(wOut)).getArray();		
			double [][] roc = ROCCurve(-0.5, 0.1, outs, dOuts);
			pw.println(Arrays.toString(roc[0]));
			pw.println(Arrays.toString(roc[1]));
			pw.println();
			System.out.print("]");
			System.out.println();
			pw.close();
			
			int score = 0;
			for(int i = 0; i < outs.length; i++){
				int v = outs[i][0] > 0.5 ? 1 : 0;
				if(v == dOuts [i][0]){
					score++;
				}
			}
			System.out.println("Accuracy: " + (double)score/outs.length);
			
		} catch (IOException e) {
			e.printStackTrace();
			System.exit(2);
		}
	}
	
	/**
	 * Creates an ROC curve from the data given a starting seed (used for
	 * finding the highest threshold for which true positives and false
	 * alarms are 100%) and an interval.
	 * @param seed
	 * @param interval
	 * @param outs
	 * @param targetOuts
	 * @return
	 */
	public double [][] ROCCurve(double seed, double interval,
			double [][] outs, double [][] targetOuts) {
		double [] tpfa;
		double tStart = findROCStart(seed, interval, outs, targetOuts);
		ArrayList<double []> roc = new ArrayList<double []>();
		double threshold = tStart;
		do {
			tpfa = ROC(threshold, outs, targetOuts);
			roc.add(Arrays.copyOf(tpfa, 2));
			threshold += interval;
		}while(tpfa[0] > 0 || tpfa[1] > 0);
		
		double [][] curve = new double[2][roc.size()];
		int i = 0;
		double auc = 0;
		for(double [] d : roc){
			curve[0][i] = d[0];
			curve[1][i] = d[1];
			if(i > 0){
				auc += (curve[1][i-1]-curve[1][i]) * curve[0][i-1];
			}
			i++;
		}
		System.out.print(auc + " ");
		return curve;
	}
	
	/**
	 * Finds the highest threshold for which true positives and false alarms
	 * are both 100%. 
	 * @param seed
	 * @param interval
	 * @param outs
	 * @param targetOuts
	 * @return
	 */
	private double findROCStart(double seed, double interval, double[][] outs,
			double [][] targetOuts) {
		double [] tpfa = ROC(seed, outs, targetOuts);
		double threshold = seed;
		if(tpfa[0] == 1 && tpfa[1] == 1) {
			do {
				threshold += interval;
				tpfa = ROC(threshold, outs, targetOuts);		
			}while(tpfa[0] == 1 && tpfa[1] == 1); 
			return threshold - interval;
		} else {
			do {
				threshold -= interval;
				tpfa = ROC(threshold, outs, targetOuts);		
			}while(tpfa[0] < 1 || tpfa[1] < 1); 
			return threshold;
		}		
	}
	
	/**
	 * Finds the ROC (Receiver Operator Characteristic: false alarm % vs.
	 * true positive %) for a given threshold given some outputs and desired
	 * outputs. 
	 * @param threshold
	 * @param outs
	 * @param targetOuts
	 * @return
	 */
	private double [] ROC(double threshold, double [][] outs,
			double [][] targetOuts) {
		
		double tp = 0;
		double fa = 0;	
		double tpMax = 0;
		double faMax = 0;
		for(int i = 0; i < outs.length; i++){
			for(int j = 0; j < outs[0].length; j++) {
				int val = outs[i][j] >= threshold ? 1 : 0;
				if(targetOuts[i][j] == 1) {
					if(val == 1) 
						tp++;
					tpMax++;
				} else {
					if(val == 1)
						fa++;
					faMax++;
				}
			}
		}
		
		return new double [] {tp/tpMax, fa/faMax};
		
		
		
	}
	
	public static void printToFile(String filename, double [][] mat) {
		try {
			FileWriter fw = new FileWriter(filename);
			PrintWriter pw = new PrintWriter(fw);
			Matrix m = new Matrix(mat);
			
			m.print(pw, 3, 4);
			
			pw.close();
			
		} catch (IOException ie) {
			ie.printStackTrace();
		}
	}
	
	
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		
		AdderTest test = new AdderTest();
		test.generateInputSets(Gate.FA_SUM);
		
		
		try{
		//First Pass
		run(test, trainingIn, false, false);
		} catch(RuntimeException e) {
			System.err.println(e.getMessage());
			double mean = 0;
			int c = 0;
			for(int i = 0; i < numRes; i++) {
				mean += test.getRes()[i].getMemPotential();
				if(test.getRes()[i].getMemPotential() > 1000 ||
						test.getRes()[i].getMemPotential() < -1000) {
					c++;
				}
			}
			mean = mean / numRes;
			System.out.println("Mean V: " + mean);
			System.out.println("Count: " + c);
			int indx = 0;
			double [][] syns = new double[1024][1024];
			for(MLNeuron n : test.getRes()){	
				for(Synapse s: n.getFanIn()) {
					syns[indx][s.getSrc().getIndex()] = s.getA();
				}
				indx++;
			}
			printToFile("SynapsesAfter.dat", syns);
			double [][] mps = new double[1][1024];
			for(MLNeuron n : test.getRes()) {
				mps[0][n.getIndex()] = n.getMemPotential();
			}
			printToFile("MemPotentials.dat", mps);
			System.exit(1);
		}
		
		for(MLNeuron n : test.getRes()) {
			n.setReporter(true);
			n.setStdpOn(false);
		}
		
		//Train Normal
		run(test, trainingIn, true, false);
		
		//Test Normal
		testFileName = "ROCsStandardRun.dat";
		run(test, testingIn, true, true);
//		
//		//Train Linearly Combined
//		runSeparate(test, trainingIn, true, false);
//		
//		//Test Linearly Combined
//		testFileName = "ROCsLCRun.dat";
//		runSeparate(test, testingIn, true, true);
		
//		run(test, trainingIn, false);
//		run(test, testingIn, true);
		

		
	}

	
	
	/*/*******************************************************************
	 * 						GETTERS AND SETTERS							 *
	 *********************************************************************/
	
	public StateCamera getCamera() {
		return camera;
	}

	public void setCamera(StateCamera camera) {
		this.camera = camera;
	}

	public static int[] getResDims() {
		return resDims;
	}

	public static void setResDims(int[] resDims) {
		AdderTest.resDims = resDims;
	}

	public boolean [][] getTrainingIn() {
		return trainingIn;
	}

	public boolean [][] getTestingIn() {
		return testingIn;
	}

	public MLNeuron[] getRes() {
		return res;
	}

	public void setRes(MLNeuron[] res) {
		this.res = res;
	}

	public MLNeuron[] getStimGrp1() {
		return stimGrp1;
	}

	public void setStimGrp1(MLNeuron[] stimGrp1) {
		this.stimGrp1 = stimGrp1;
	}

	public MLNeuron[] getStimGrp2() {
		return stimGrp2;
	}

	public void setStimGrp2(MLNeuron[] stimGrp2) {
		this.stimGrp2 = stimGrp2;
	}

	public MLNeuron[] getStimGrp3() {
		return stimGrp3;
	}

	public void setStimGrp3(MLNeuron[] stimGrp3) {
		this.stimGrp3 = stimGrp3;
	}

	public MLNeuron [] getStimGrp(int grpNum){
		switch(grpNum){
			case 1: return getStimGrp1();
			case 2: return getStimGrp2();
			case 3: return getStimGrp3();
			default: throw new IllegalArgumentException("No such stim" +
					" group.");
		}
	}

	public static int getTestepochs() {
		return testEpochs;
	}

	public static double getTimestep() {
		return timeStep;
	}

	public static double getEpoch() {
		return epoch;
	}

	public static int getNumres() {
		return numRes;
	}

	public static int getMemtc() {
		return memTc;
	}

	public Network getNet() {
		return net;
	}
}
